package com.vrv.utils;

import com.vrv.entity.ElasticsearchProperties;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.StatusLine;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.Credentials;
import org.apache.http.auth.UsernamePasswordCredentials;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.entity.ContentType;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.nio.entity.NStringEntity;
import org.apache.http.util.EntityUtils;
import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
import org.elasticsearch.action.admin.indices.get.GetIndexRequest;
import org.elasticsearch.action.search.ClearScrollRequest;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchScrollRequest;
import org.elasticsearch.action.support.master.AcknowledgedResponse;
import org.elasticsearch.client.*;
import org.elasticsearch.core.TimeValue;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.time.Duration;
import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * @author vrv
 * @ClassName: EsUtils
 * @Desc: TODO
 * @Date: 2024/07/04 16:11
 * @Version: 1.0
 * @Modified By:
 */
public class EsUtils {


    public static String getAllTemplate(RestHighLevelClient client) throws IOException {
        RestClient levelClient = client.getLowLevelClient();
        Request request = new Request("GET", "/_template/*" );
        Response response = levelClient.performRequest(request);
        StatusLine statusLine = response.getStatusLine();
        if (statusLine.getStatusCode() == 200) {
            String string = EntityUtils.toString(response.getEntity());
            return string;
        }
        return null;
    }

    public static void putTemplate(RestHighLevelClient client, String name, String templateJson) throws IOException {
        HttpEntity entity = new NStringEntity(templateJson, ContentType.APPLICATION_JSON);
        RestClient levelClient = client.getLowLevelClient();
        Request request = new Request("PUT", "/_template/" + name);
        request.setEntity(entity);
        levelClient.performRequest(request);
    }

    public static boolean existsTemplate(RestHighLevelClient client, String name) {
        RestClient levelClient = client.getLowLevelClient();
        final Response response;
        try {
            Request request = new Request("HEAD", "/_template/" + name);
            response = levelClient.performRequest(request);
            StatusLine statusLine = response.getStatusLine();
            if (statusLine.getStatusCode() == 200) {
                return true;
            }
        } catch (IOException e) {
            return false;
        }
        return false;

    }

    /**
     * scrollSearch
     *
     * @param client
     * @param index
     * @return
     * @throws IOException
     */
    public static boolean exist(RestHighLevelClient client, String index) throws IOException {
        return client.indices().exists(new GetIndexRequest().indices(index),RequestOptions.DEFAULT);
    }

    public static boolean deleteIndex(RestHighLevelClient client, String index) throws IOException {
        boolean result = false;
        final AcknowledgedResponse delete;
        try {
            delete = client.indices().delete(new DeleteIndexRequest(index),RequestOptions.DEFAULT);
            result = delete.isAcknowledged();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return result;
    }

    private static final Pattern SUCCESS_PATTERN = Pattern.compile("20\\d+");


    /**
     * scrollSearch
     *
     * @param client
     * @param index
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> scrollSearch(RestHighLevelClient client, Integer size, String... index) throws IOException {
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(size);
        return scrollSearch(client, searchSourceBuilder, index);
    }

    /**
     * scrollSearch
     *
     * @param client
     * @param searchSourceBuilder
     * @param index
     * @return
     * @throws IOException
     */
    public static List<Map<String, Object>> scrollSearch(RestHighLevelClient client, SearchSourceBuilder searchSourceBuilder, String... index) throws IOException {
        final Scroll scroll = new Scroll(TimeValue.timeValueMinutes(30L));
        SearchRequest searchRequest = new SearchRequest(index);
        searchRequest.scroll(scroll);
        searchRequest.source(searchSourceBuilder);
        SearchResponse searchResponse = client.search(searchRequest, RequestOptions.DEFAULT);
        String scrollId = searchResponse.getScrollId();
        SearchHit[] searchHits = searchResponse.getHits().getHits();
        List<Map<String, Object>> list = new ArrayList<>();
        while (searchHits != null && searchHits.length > 0) {
            for (int i = 0; i < searchHits.length; i++) {
                SearchHit searchHit = searchHits[i];
                if (searchHit.getSourceAsMap() != null) {
                    list.add(searchHit.getSourceAsMap());
                }
            }
            SearchScrollRequest scrollRequest = new SearchScrollRequest(scrollId);
            scrollRequest.scroll(scroll);
            searchResponse = client.searchScroll(scrollRequest, RequestOptions.DEFAULT);
            scrollId = searchResponse.getScrollId();
            searchHits = searchResponse.getHits().getHits();
        }

        ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
        clearScrollRequest.addScrollId(scrollId);
        client.clearScroll(clearScrollRequest, RequestOptions.DEFAULT);
        return list;
    }


    /**
     * 默认返回 connectionTimeout:30s readTimeout:30s
     *
     * @param properties
     * @return
     */
    public static RestHighLevelClient getRestHighLevelClient(ElasticsearchProperties properties) {
        HttpHost[] hosts = properties.getUris().stream().map(EsUtils::createHttpHost).toArray(HttpHost[]::new);
        RestClientBuilder builder = RestClient.builder(hosts);

        final BasicCredentialsProvider provider = new BasicCredentialsProvider();
        if (StringUtils.hasText(properties.getUsername())) {
            Credentials credentials = new UsernamePasswordCredentials(properties.getUsername(),
                    properties.getPassword());
            provider.setCredentials(AuthScope.ANY, credentials);
        }
        properties.getUris().stream().map(EsUtils::toUri).filter(EsUtils::hasUserInfo)
                .forEach(e -> addUserInfoCredentials(e, provider));

        builder.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder.setDefaultCredentialsProvider(provider));
        builder.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
            @Override
            public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
                Duration connectionTimeout = properties.getConnectionTimeout();
                if (Objects.isNull(connectionTimeout)) {
                    connectionTimeout = Duration.ofSeconds(30L);
                }
                Duration readTimeout = properties.getReadTimeout();
                if (Objects.isNull(readTimeout)) {
                    readTimeout = Duration.ofSeconds(30);
                }
                return requestConfigBuilder.setConnectTimeout(Math.toIntExact(connectionTimeout.toMillis())).setSocketTimeout(Math.toIntExact(readTimeout.toMillis()));
            }
        });
        return new RestHighLevelClient(builder);
    }

    private static URI toUri(String uri) {
        try {
            return URI.create(uri);
        } catch (IllegalArgumentException ex) {
            return null;
        }
    }

    private static boolean hasUserInfo(URI uri) {
        return uri != null && StringUtils.hasText(uri.getUserInfo());
    }

    private static void addUserInfoCredentials(URI uri, BasicCredentialsProvider provider) {
        AuthScope authScope = new AuthScope(uri.getHost(), uri.getPort());
        Credentials credentials = createUserInfoCredentials(uri.getUserInfo());
        provider.setCredentials(authScope, credentials);
    }

    private static Credentials createUserInfoCredentials(String userInfo) {
        int delimiter = userInfo.indexOf(":");
        if (delimiter == -1) {
            return new UsernamePasswordCredentials(userInfo, null);
        }
        String username = userInfo.substring(0, delimiter);
        String password = userInfo.substring(delimiter + 1);
        return new UsernamePasswordCredentials(username, password);
    }

    private static HttpHost createHttpHost(String uri) {
        try {
            return createHttpHost(URI.create(uri));
        } catch (IllegalArgumentException ex) {
            return HttpHost.create(uri);
        }
    }

    private static HttpHost createHttpHost(URI uri) {
        if (!StringUtils.hasText(uri.getUserInfo())) {
            return HttpHost.create(uri.toString());
        }
        try {
            return HttpHost.create(new URI(uri.getScheme(), null, uri.getHost(), uri.getPort(), uri.getPath(),
                    uri.getQuery(), uri.getFragment()).toString());
        } catch (URISyntaxException ex) {
            throw new IllegalStateException(ex);
        }
    }

    public static List<HttpHost> getHttpHosts(ElasticsearchProperties properties) {
        return properties.getUris().stream().map(EsUtils::createHttpHost).collect(Collectors.toList());
    }

}
